"""
记忆工厂类 - 整合Mem0和传统记忆管理
支持多种记忆类型：LangChain记忆、Mem0记忆、数据库记忆
"""

import uuid
from datetime import datetime
from typing import Dict, Any, Optional, Union

from langchain.memory import (
    ConversationBufferMemory,
    ConversationBufferWindowMemory,
    ConversationSummaryMemory,
    ConversationSummaryBufferMemory,
    ConversationTokenBufferMemory
)
from langchain.schema import BaseMemory
from sqlalchemy.orm import Session

from config.logging_config import get_logger
from .Mem0Interface import Mem0Interface, Mem0MemoryProvider
from .MemoryManager import MemoryManager


logging = get_logger(__name__)

class MemoryFactory:
    """增强的记忆工厂类，支持多种记忆类型"""

    # LangChain记忆类型映射
    _langchain_memory_types = {
        "buffer": ConversationBufferMemory,
        "buffer_window": ConversationBufferWindowMemory,
        "summary": ConversationSummaryMemory,
        "summary_buffer": ConversationSummaryBufferMemory,
        "token_buffer": ConversationTokenBufferMemory
    }

    @classmethod
    def create_memory_system(
            cls,
            memory_config: Dict[str, Any],
            db_session: Session,
            scene_id: str,
            agent_id: str,
            session_id: Optional[str] = None,
            user_id: str = "1"
    ) -> 'MemoryManager':
        """
        创建完整的记忆管理系统
        
        Args:
            memory_config: 记忆配置
            db_session: 数据库会话
            scene_id: 场景ID
            agent_id: 智能体ID
            session_id: 会话ID（可选，会自动生成）
            user_id: 用户ID
            
        Returns:
            MemoryManager实例
        """
        logger = get_logger(__name__)

        try:
            # 生成session_id（如果没有提供）
            if not session_id:
                session_id = f"session_{uuid.uuid4().hex[:8]}_{int(datetime.now().timestamp())}"

            # 获取记忆类型，默认为mem0
            memory_type = memory_config.get("type", "mem0")
            # 初始化组件
            langchain_memory = None
            mem0_provider = None

            # 1. 尝试创建Mem0记忆（优先）
            if memory_type == "mem0" or memory_config.get("enable_mem0", True):
                try:
                    mem0_interface = Mem0Interface()
                    if mem0_interface.is_available():
                        mem0_provider = Mem0MemoryProvider(mem0_interface, db_session)
                        logger.info("✅ Mem0记忆接口创建成功")
                    else:
                        error_msg = mem0_interface.get_initialization_error()
                        logger.warning(f"⚠️ Mem0记忆接口创建失败: {error_msg}")
                except Exception as e:
                    logger.warning(f"⚠️ Mem0记忆初始化失败: {e}")

            # 2. 创建LangChain记忆（作为补充或主要记忆）
            if memory_type in cls._langchain_memory_types or not mem0_provider:
                try:
                    langchain_type = memory_type if memory_type in cls._langchain_memory_types else "buffer"
                    langchain_params = memory_config.get("langchain_params", {})
                    langchain_memory = cls._create_langchain_memory(langchain_type, **langchain_params)
                    if langchain_memory:
                        logger.info(f"✅ LangChain记忆创建成功: {langchain_type}")
                except Exception as e:
                    logger.warning(f"⚠️ LangChain记忆创建失败: {e}")

            # 3. 创建记忆管理器
            memory_manager = MemoryManager(
                langchain_memory=langchain_memory,
                mem0_provider=mem0_provider,
                db_session=db_session,
                scene_id=scene_id,
                agent_id=agent_id,
                session_id=session_id,
                user_id=user_id,
                config=memory_config
            )

            # 4. 预加载历史记忆（可选）
            if memory_config.get("preload_history", False):
                try:
                    memory_manager.preload_session_history()
                    logger.info(f"✅ 会话历史预加载完成")
                except Exception as e:
                    logger.warning(f"⚠️ 预加载会话历史失败: {e}")
            return memory_manager

        except Exception as e:
            logger.error(f"❌ 创建记忆系统失败: {e}")
            # 返回最基础的记忆管理器
            try:
                basic_memory = cls._create_langchain_memory("buffer")
                return MemoryManager(
                    langchain_memory=basic_memory,
                    mem0_provider=None,
                    db_session=db_session,
                    scene_id=scene_id,
                    agent_id=agent_id,
                    session_id=session_id or f"fallback_{uuid.uuid4().hex[:8]}",
                    user_id=user_id,
                    config={"type": "buffer"}
                )
            except Exception as fallback_error:
                logger.error(f"❌ 创建备用记忆系统也失败: {fallback_error}")
                raise

    @classmethod
    def _create_langchain_memory(cls, memory_type: str, **kwargs) -> Optional[BaseMemory]:
        """创建LangChain记忆实例"""
        try:
            memory_class = cls._langchain_memory_types.get(memory_type)
            if memory_class:
                # 设置默认参数
                default_params = {
                    "return_messages": True,
                    "memory_key": "chat_history",
                    "input_key": "input"
                }

                # 合并用户参数
                params = {**default_params, **kwargs}

                # 特殊处理不同类型的记忆
                if memory_type == "buffer_window":
                    params.setdefault("k", 10)  # 默认保留最近10轮对话
                elif memory_type == "token_buffer":
                    params.setdefault("max_token_limit", 2000)  # 默认2000 token限制
                elif memory_type in ["summary", "summary_buffer"]:
                    # 摘要类型记忆需要LLM，暂时跳过
                    logging.warning(f"摘要类型记忆 {memory_type} 需要LLM配置，使用buffer替代")
                    return cls._create_langchain_memory("buffer", **kwargs)

                return memory_class(**params)

        except Exception as e:
            logging.error(f"创建LangChain记忆失败: {e}")

        return None

    @classmethod
    def create_from_config(
            cls,
            config: Dict[str, Any],
            db_session: Session = None,
            scene_id: str = None,
            agent_id: str = None,
            session_id: Optional[str] = None,
            user_id: str = "1",
            agent: Any = None,
            **kwargs
    ) -> Optional[Union[BaseMemory, 'MemoryManager']]:
        """
        从配置创建记忆系统
        
        Args:
            config: 记忆配置
            db_session: 数据库会话  
            scene_id: 场景ID
            agent_id: 智能体ID
            session_id: 会话ID（可选）
            user_id: 用户ID
            agent: 智能体实例（用于设置记忆管理器）
            **kwargs: 额外参数
            
        Returns:
            MemoryManager实例或BaseMemory实例（向后兼容）
        """
        logger = get_logger(__name__)

        try:
            # 验证配置
            validated_config = cls.validate_config(config)
            
            # 如果提供了完整的参数，创建完整的记忆系统
            if db_session and scene_id and agent_id:
                memory_manager = cls.create_memory_system(
                    memory_config=validated_config,
                    db_session=db_session,
                    scene_id=scene_id,
                    agent_id=agent_id,
                    session_id=session_id,
                    user_id=user_id
                )

                # 设置到Agent
                if agent and memory_manager:
                    cls._set_memory_to_agent(agent, memory_manager, logger)

                return memory_manager

            else:
                # 向后兼容：只创建LangChain记忆
                logger.info("创建简单LangChain记忆（向后兼容）")
                memory_type = validated_config.get("type", "buffer")
                params = validated_config.get("langchain_params", {})
                langchain_memory = cls._create_langchain_memory(memory_type, **params)

                # 设置到Agent
                if agent and langchain_memory:
                    if hasattr(agent, 'memory'):
                        agent.memory = langchain_memory
                        logger.info("✅ Agent LangChain记忆设置成功（兼容模式）")

                return langchain_memory

        except Exception as e:
            logger.error(f"❌ 从配置创建记忆失败: {e}")
            # 返回基础记忆作为备用
            try:
                return cls._create_langchain_memory("buffer")
            except Exception as fallback_error:
                logger.error(f"❌ 创建备用记忆也失败: {fallback_error}")
                return None

    @classmethod
    def _set_memory_to_agent(cls, agent: Any, memory_manager: 'MemoryManager', logger) -> bool:
        """设置记忆管理器到Agent"""
        try:
            if hasattr(agent, 'set_memory_manager'):
                agent.set_memory_manager(memory_manager)
                logger.info("✅ Agent记忆管理器设置成功")
                return True
            elif hasattr(agent, 'memory_manager'):
                agent.memory_manager = memory_manager
                logger.info("✅ Agent记忆管理器属性设置成功")
                return True
            else:
                logger.warning("⚠️ Agent不支持记忆管理器设置")
                return False
        except Exception as e:
            logger.error(f"❌ 设置记忆到Agent失败: {e}")
            return False

    @classmethod
    def get_available_types(cls) -> Dict[str, str]:
        """获取可用的记忆类型"""
        return {
            "mem0": "Mem0智能记忆",
            "buffer": "缓冲记忆",
            "buffer_window": "窗口缓冲记忆",
            "summary": "摘要记忆",
            "summary_buffer": "摘要缓冲记忆",
            "token_buffer": "令牌缓冲记忆"
        }

    @classmethod
    def validate_config(cls, config: Dict[str, Any]) -> Dict[str, Any]:
        """
        验证和规范化记忆配置
        
        Args:
            config: 原始配置
            
        Returns:
            验证后的配置
        """
        validated_config = config.copy() if config else {}

        # 验证记忆类型，默认使用 mem0
        memory_type = validated_config.get("type", "mem0")
        available_types = list(cls.get_available_types().keys())

        if memory_type not in available_types:
            logging.warning(f"不支持的记忆类型 {memory_type}，使用默认 mem0")
            validated_config["type"] = "mem0"

        # 设置默认参数（简化）
        defaults = {
            "preload_history": False,  # 默认不预加载，提高启动速度
            "auto_save": True,
            "enable_mem0": True,
            "langchain_params": {
                "return_messages": True,
                "memory_key": "chat_history",
                "input_key": "input"
            }
        }

        for key, default_value in defaults.items():
            if key not in validated_config:
                validated_config[key] = default_value

        return validated_config

    @classmethod
    def create_session_memory(
            cls,
            session_id: str,
            scene_id: str,
            agent_id: str,
            user_id: str = "1",
            memory_type: str = "mem0",
            db_session: Optional[Session] = None
    ) -> Optional['MemoryManager']:
        """
        快速创建会话记忆
        
        Args:
            session_id: 会话ID
            scene_id: 场景ID
            agent_id: 智能体ID
            user_id: 用户ID
            memory_type: 记忆类型
            db_session: 数据库会话
            
        Returns:
            MemoryManager实例
        """
        if not db_session:
            logging.error("创建会话记忆需要数据库会话")
            return None

        config = {
            "type": memory_type,
            "preload_history": True,
            "auto_save": True
        }

        return cls.create_memory_system(
            memory_config=config,
            db_session=db_session,
            scene_id=scene_id,
            agent_id=agent_id,
            session_id=session_id,
            user_id=user_id
        )

    @classmethod
    def get_memory_info(cls, memory: BaseMemory) -> Dict[str, Any]:
        """
        获取记忆信息
        
        Args:
            memory: 记忆实例
            
        Returns:
            记忆信息字典
        """
        if not memory:
            return {"type": "none", "status": "未初始化"}

        memory_type = type(memory).__name__
        info = {
            "type": memory_type,
            "status": "已初始化",
            "memory_key": getattr(memory, 'memory_key', 'unknown'),
            "return_messages": getattr(memory, 'return_messages', False)
        }

        # 获取特定类型的额外信息
        if hasattr(memory, 'chat_memory') and hasattr(memory.chat_memory, 'messages'):
            info["message_count"] = len(memory.chat_memory.messages)

        if hasattr(memory, 'k'):
            info["window_size"] = memory.k

        if hasattr(memory, 'max_token_limit'):
            info["max_tokens"] = memory.max_token_limit

        return info
